import os
import torch
import numpy as np
import cv2
from pathlib import Path
import tkinter as tk
from tkinter import filedialog, messagebox, ttk
from PIL import Image, ImageTk
import time

# 获取当前脚本所在目录
current_dir = os.path.dirname(os.path.abspath(__file__))
# 构建真实处理.txt的完整路径
real_process_path = os.path.join(current_dir, "真实处理.txt")

# 检查文件是否存在
if not os.path.exists(real_process_path):
    raise FileNotFoundError(f"无法找到文件: {real_process_path}")

# 直接执行文件内容来导入函数
namespace = {}
with open(real_process_path, 'r', encoding='utf-8') as f:
    exec(f.read(), namespace)

repair_pipeline = namespace['repair_pipeline']

def select_image():
    """选择图像文件"""
    file_path = filedialog.askopenfilename(
        title="选择牛顿环图像",
        filetypes=[("图像文件", "*.png;*.jpg;*.jpeg;*.bmp")]
    )
    if file_path:
        entry_path.delete(0, tk.END)
        entry_path.insert(0, file_path)
        # 显示图像预览
        show_preview(file_path)


def show_preview(file_path):
    """显示图像预览"""
    try:
        # 读取图像
        img = Image.open(file_path)
        # 调整大小以适应预览窗口
        img.thumbnail((300, 300))
        # 转换为PhotoImage对象
        photo = ImageTk.PhotoImage(img)
        # 更新预览标签
        label_preview.config(image=photo)
        label_preview.image = photo  # 保持引用，防止被垃圾回收
    except Exception as e:
        messagebox.showerror("错误", f"无法预览图像: {str(e)}")


def start_repair():
    """开始修复图像"""
    input_path = entry_path.get().strip()
    if not input_path:
        messagebox.showwarning("警告", "请先选择图像文件")
        return

    if not Path(input_path).exists():
        messagebox.showerror("错误", f"无法找到图像文件: {input_path}")
        return

    # 检查模型是否存在
    model_path = "results/G_epoch_50.pth"
    if not Path(model_path).exists():
        # 尝试查找results目录下的其他模型文件
        results_dir = Path("results")
        model_files = list(results_dir.glob("G_epoch_*.pth"))
        if model_files:
            # 选择最新的模型文件（按文件名排序）
            model_files.sort(key=lambda x: int(x.stem.split('_')[-1]))
            model_path = str(model_files[-1])
            messagebox.showinfo("信息", f"找到模型文件: {model_path}")
        else:
            messagebox.showerror("错误", "模型文件不存在。请先运行'牛顿环.txt'来训练模型。")
            return

    # 创建输出目录
    output_dir = "real_image_results"
    Path(output_dir).mkdir(exist_ok=True)

    # 更新状态
    label_status.config(text="正在修复图像...")
    root.update_idletasks()

    try:
        # 执行修复
        start_time = time.time()
        success = repair_pipeline(input_path, model_path, output_dir)
        end_time = time.time()

        if success:
            # 读取报告
            report_path = Path(output_dir) / "process_report.txt"
            if report_path.exists():
                with open(report_path, "r", encoding="utf-8") as f:
                    report_content = f.read()
                messagebox.showinfo("修复完成", f"修复成功! 结果保存在: {output_dir}\n耗时: {end_time - start_time:.2f}秒\n\n{report_content}")
            else:
                messagebox.showinfo("修复完成", f"修复成功! 结果保存在: {output_dir}\n耗时: {end_time - start_time:.2f}秒")
        else:
            messagebox.showerror("修复失败", "修复过程中发生错误，请查看日志文件获取详细信息。")
    except Exception as e:
        messagebox.showerror("错误", f"修复过程中发生异常: {str(e)}")
    finally:
        label_status.config(text="就绪")


# 创建GUI界面
root = tk.Tk()
root.title("牛顿环图像修复系统")
root.geometry("600x500")
root.resizable(True, True)

# 设置中文字体
font_style = ("SimHei", 10)

# 创建框架
frame_main = ttk.Frame(root, padding="20")
frame_main.pack(fill=tk.BOTH, expand=True)

# 标题
label_title = ttk.Label(frame_main, text="牛顿环图像修复与曲率半径计算", font=("SimHei", 16, "bold"))
label_title.pack(pady=10)

# 图像路径选择
frame_path = ttk.Frame(frame_main)
frame_path.pack(fill=tk.X, pady=10)

label_path = ttk.Label(frame_path, text="图像路径:", font=font_style)
label_path.pack(side=tk.LEFT, padx=5)

entry_path = ttk.Entry(frame_path, width=40, font=font_style)
entry_path.pack(side=tk.LEFT, padx=5, fill=tk.X, expand=True)

button_browse = ttk.Button(frame_path, text="浏览...", command=select_image)
button_browse.pack(side=tk.RIGHT, padx=5)

# 图像预览
frame_preview = ttk.LabelFrame(frame_main, text="图像预览")
frame_preview.pack(fill=tk.BOTH, expand=True, pady=10)

label_preview = ttk.Label(frame_preview, text="请选择图像文件")
label_preview.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)

# 操作按钮
frame_actions = ttk.Frame(frame_main)
frame_actions.pack(fill=tk.X, pady=10)

button_repair = ttk.Button(frame_actions, text="开始修复", command=start_repair)
button_repair.pack(side=tk.LEFT, padx=5)

button_exit = ttk.Button(frame_actions, text="退出", command=root.quit)
button_exit.pack(side=tk.RIGHT, padx=5)

# 状态标签
label_status = ttk.Label(frame_main, text="就绪", font=font_style)
label_status.pack(side=tk.BOTTOM, pady=10)

# 运行主循环
root.mainloop()